import cv2
import numpy as np

def calculate_mse_psnr(original_path, extracted_path):
    """
    计算均方误差 (MSE) 和峰值信噪比 (PSNR)
    """
    # 读取原始水印和提取水印图像
    original = cv2.imread(original_path, cv2.IMREAD_GRAYSCALE)
    extracted = cv2.imread(extracted_path, cv2.IMREAD_GRAYSCALE)

    # 确保图像尺寸一致
    if original.shape != extracted.shape:
        extracted = cv2.resize(extracted, (original.shape[1], original.shape[0]))

    # 计算均方误差 (MSE)
    mse = np.mean((original - extracted) ** 2)

    return mse


# 进行质量评估
if __name__ == "__main__":
    original_path = "dog.png"
    extracted_paths = [
        "extracted_from_hist.png",
        "extracted_from_smooth.png",
        "extracted_from_sharpen.png"
    ]

    for extracted_path in extracted_paths:
        mse = calculate_mse_psnr(original_path, extracted_path)
        print(f"提取结果: {extracted_path}")
        print(f"均方误差 (MSE): {mse:.2f}")
